package com.jxpanda.commons.key;

import com.jxpanda.commons.constant.SchemeConstant;
import com.jxpanda.commons.constant.StringConstant;
import com.jxpanda.commons.toolkit.Base64Kit;
import com.jxpanda.commons.encrypt.EncryptAlgorithm;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

/**
 * 密钥对工具
 * 生成密钥对的算法不在这里
 *
 * @author Panda
 */
@Slf4j
public class KeyHelper {

    private static final String REGEX_SPACE = "\\s+";

    private static final String PUBLIC_KEY_PREFIX = "-----BEGIN PUBLIC KEY-----";
    private static final String PUBLIC_KEY_SUFFIX = "-----END PUBLIC KEY-----";

    private static final String PRIVATE_KEY_PREFIX = "-----BEGIN PRIVATE KEY-----";
    private static final String PRIVATE_KEY_SUFFIX = "-----END PRIVATE KEY-----";

    protected String algorithm;
    protected KeyFactory keyFactory;
    protected KeyPairGenerator keyPairGenerator;

    public KeyHelper(EncryptAlgorithm algorithm, int keySize) {
        this.algorithm = algorithm.getValue();
        try {
            this.keyFactory = KeyFactory.getInstance(this.algorithm);
            this.keyPairGenerator = KeyPairGenerator.getInstance(this.algorithm);
            this.keyPairGenerator.initialize(keySize);
        } catch (NoSuchAlgorithmException e) {
            log.error("【NO SUCH ALGORITHM】KeyHelper 创建失败，不支持的算法。", e);
        }
    }

    protected String trimPublic(String publicKey) {
        return publicKey.replace(PUBLIC_KEY_PREFIX, StringConstant.BLANK)
                .replace(PUBLIC_KEY_SUFFIX, StringConstant.BLANK)
                .replaceAll(REGEX_SPACE, StringConstant.BLANK);
    }

    protected String trimPrivate(String privateKey) {
        return privateKey.replace(PRIVATE_KEY_PREFIX, StringConstant.BLANK)
                .replace(PRIVATE_KEY_SUFFIX, StringConstant.BLANK)
                .replaceAll(REGEX_SPACE, StringConstant.BLANK);
    }

    /**
     * 理论上不会频繁生成密钥对，所以不需要缓存KeyPairGenerator对象
     * 因此使用局部变量就可以了
     */
    @SneakyThrows
    public KeyPair generateKeyPair() {
        return keyPairGenerator.generateKeyPair();
    }

    /**
     * byte数组转公钥
     */
    @SneakyThrows
    public PublicKey toPublicKey(byte[] bytes) {
        return keyFactory.generatePublic(new X509EncodedKeySpec(bytes));
    }

    /**
     * 字符串转公钥
     */
    public PublicKey toPublicKey(String publicKey) {
        return toPublicKey(Base64Kit.decode(trimPublic(publicKey)));
    }

    /**
     * InputStream转公钥
     * IO异常都往外抛出，谁调用，谁处理
     */
    public PublicKey toPublicKey(InputStream inputStream) throws IOException {
        return toPublicKey(new String(inputStream.readAllBytes(), StandardCharsets.UTF_8));
    }


    /**
     * 公钥对象转字符串
     */
    public String fromPublicKey(PublicKey publicKey) {
        return Base64Kit.encodeToString(publicKey.getEncoded());
    }

    /**
     * byte数组转私钥
     */
    @SneakyThrows
    public PrivateKey toPrivateKey(byte[] bytes) {
        return keyFactory.generatePrivate(new PKCS8EncodedKeySpec(bytes));
    }


    /**
     * 字符串转私钥
     */
    public PrivateKey toPrivateKey(String privateKey) {
        return toPrivateKey(Base64Kit.decode(trimPrivate(privateKey)));
    }

    /**
     * InputStream转私钥
     * IO异常都往外抛出，谁调用，谁处理
     */
    public PrivateKey toPrivateKey(InputStream inputStream) throws IOException {
        return toPrivateKey(new String(inputStream.readAllBytes(), StandardCharsets.UTF_8));
    }


    public String fromPrivateKey(PrivateKey privateKey) {
        return Base64Kit.encodeToString(privateKey.getEncoded());
    }

    public PublicKey loadPublicKey(String publicKey) {
        if (isNotFile(publicKey)) {
            return toPublicKey(publicKey);
        } else {
            try (FileInputStream inputStream = loadFile(publicKey)) {
                return toPublicKey(inputStream);
            } catch (Exception ex) {
                log.error("【PUBLIC KEY LOAD ERROR】读取公钥失败。", ex);
            }
        }
        return null;
    }

    public PrivateKey loadPrivateKey(String privateKey) {
        if (isNotFile(privateKey)) {
            return toPrivateKey(privateKey);
        } else {
            try (FileInputStream inputStream = loadFile(privateKey)) {
                return toPrivateKey(inputStream);
            } catch (Exception ex) {
                log.error("【PRIVATE KEY LOAD ERROR】读取私钥失败。", ex);
            }
        }
        return null;
    }

    protected static boolean isNotFile(String key) {
        return !key.startsWith(SchemeConstant.FILE);
    }

    /**
     * 读取文件，异常往外抛，谁调用，谁处理
     */
    protected static FileInputStream loadFile(String path) throws FileNotFoundException {
        return new FileInputStream(path.replace(SchemeConstant.FILE, StringConstant.BLANK));
    }

}